#include "SectionCalculator.h"

#include <cassert>
#include <iostream>

#include <Eigen/Sparse>

#include "SectionElement/ElementFactory.h"

using SG::Algebra::Matrixd;
using SG::DataStructure::Common::Real;
using SG::DataStructure::FEM::BeamMeshedSectionProperty;
using SG::FEMSolver::BeamElement::BeamSecElementFactory;  // 工厂类
using SG::FEMSolver::BeamElement::BeamSecElementTypes;
using SG::FEMSolver::BeamElement::EleGeomProp;
using SG::FEMSolver::BeamElement::ISectionElement;
using SG::FEMSolver::Element::SectionCalculator;

using EigenVectorXR = Eigen::VectorX<Real>;
using EigenSparseR  = Eigen::SparseMatrix<Real>;
using EigenTripletR = Eigen::Triplet<Real>;

void SectionCalculator::Calculate (Meshes& meshes, _OUT BeamMeshedSectionProperty& sectionProperty)
{
    // 计算面积 惯性矩 形心坐标
    CalculateNonwarpPop (meshes, _OUT sectionProperty);
    // 转换坐标到形心坐标系
    meshes.TranslateNodes (sectionProperty.m_yC, sectionProperty.m_zC);
    // 截面翘曲相关属性计算
    SectionCalculator::CalculateWarpPop (meshes, _OUT sectionProperty);
}

void SectionCalculator::CalculateNonwarpPop (const Meshes& meshes, _OUT BeamMeshedSectionProperty& sectionProperty)
{
    // 初始化,防止重复调用多次相加。
    sectionProperty.m_area = 0.0;
    sectionProperty.m_Iy   = 0.0;
    sectionProperty.m_Iz   = 0.0;
    sectionProperty.m_Iyz  = 0.0;
    sectionProperty.m_yC   = 0.0;
    sectionProperty.m_zC   = 0.0;
    Real Qy{ 0.0 };
    Real Qz{ 0.0 };

    // 从Mesh获取当前存在的单元类型
    const auto existingTypes = meshes.GetAllElementTypes ();
    if (existingTypes.empty ())
    {
        SGFatal << "No element data exists in the mesh!" << std::endl;
        return;
    }

    for (const auto& eleType : existingTypes)
    {
        auto        eleCalc = BeamSecElementFactory::Create (eleType);
        const auto& eleSet  = meshes.GetElement (eleType);
        if (eleSet.empty()) continue;
        const std::size_t nodeCount = eleSet[0].nodes.size();  // 单元节点数
        Matrixd           nodeCoords(nodeCount, 2);            // 单元节点坐标
        EleGeomProp eleGeomProp;                               // 单元集合属性
        for (const auto& ele : eleSet)
        {
            nodeCoords = meshes.ExtractElementNodeCoords (ele.nodes);  // 获取单元节点坐标
            eleCalc->CalculateGeomProp (nodeCoords);
            eleGeomProp = eleCalc->GetGeomProp ();
            sectionProperty.m_area += eleGeomProp.Area;
            Qy += eleGeomProp.Qy;
            Qz += eleGeomProp.Qz;
            sectionProperty.m_Iy += eleGeomProp.Iy;
            sectionProperty.m_Iz += eleGeomProp.Iz;
            sectionProperty.m_Iyz += eleGeomProp.Iyz;
        }
    }
    if (sectionProperty.m_area < 1.0e-6)
    {
        SGFatal << "Section area is near zero!" << std::endl;
        return;
    }

    // 计算截面形心坐标
    sectionProperty.m_yC = Qz / sectionProperty.m_area;
    sectionProperty.m_zC = Qy / sectionProperty.m_area;

    // 调整截面惯性矩（对形心轴）
    sectionProperty.m_Iy -= sectionProperty.m_area * sectionProperty.m_zC * sectionProperty.m_zC;
    sectionProperty.m_Iz -= sectionProperty.m_area * sectionProperty.m_yC * sectionProperty.m_yC;
    sectionProperty.m_Iyz -= sectionProperty.m_area * sectionProperty.m_yC * sectionProperty.m_zC;
}

// 计算翘曲相关截面属性，扭转常数、剪心坐标、剪切系数
void SectionCalculator::CalculateWarpPop (const Meshes&                   meshes,  // 传入的坐标应该为形心轴坐标
                                          _OUT BeamMeshedSectionProperty& sectionProperty)
{
    // 初始化
    std::size_t                totalDof = meshes.nodes.size ();
    Real                       Iy{ sectionProperty.m_Iy };
    Real                       Iz{ sectionProperty.m_Iz };
    Real                       Iyz{ sectionProperty.m_Iyz };
    Matrixd                    Ke;
    std::vector<Matrixd>       FSet (3);
    std::vector<EigenTripletR> kTriplet;
    EigenSparseR               K (totalDof, totalDof);
    EigenVectorXR              Pw (totalDof);   // 翘曲函数有限元方程荷载项
    EigenVectorXR              PVy (totalDof);  // Vy荷载下应力函数有限元方程荷载项
    EigenVectorXR              PVz (totalDof);  // Vz荷载下应力函数有限元方程荷载项
    Pw.setZero ();
    PVy.setZero ();
    PVz.setZero ();
    const Real mu{ 0.0 };  // 这里暂时把泊松比初始化为0，后期若有需要再修改。

    //***********************************************组装整体刚度矩阵和荷载向量***********************************************************
    // 从Mesh获取当前存在的单元类型
    const auto existingTypes = meshes.GetAllElementTypes ();
    if (existingTypes.empty ())
    {
        SGFatal << "No element data exists in the mesh!" << std::endl;
        return;
    }

    for (const auto& type : existingTypes)
    {
        auto elementCalc = BeamSecElementFactory::Create (type);
        assert (elementCalc && "Failed to create element calculator");
        const auto& eleSet = meshes.GetElement (type);
        if (eleSet.empty()) continue;                            // 补充防御性判断，避免eleSet[0]越界
        const std::size_t nodeCount = eleSet[0].nodes.size();    // 单元节点个数
        Matrixd nodeCoords(nodeCount, 2);                        // 单元节点坐标,提前定义避免多次析构
        Matrixd eleKe(nodeCount, nodeCount);                     // 单元刚度矩阵
        std::vector<SG::Algebra::Matrixd> ele_FSet(3);  // 单元荷载向量集
        
        for (const auto& element : eleSet)
        {
            nodeCoords = meshes.ExtractElementNodeCoords (element.nodes);
            // 计算单元刚度矩阵和荷载项
            elementCalc->ComputeEleMatrix(nodeCoords, {Iy, Iz, Iyz}, mu, eleKe, ele_FSet);
            // 组装单元刚度和荷载项
            for (std::size_t iNode = 0; iNode < nodeCount; ++iNode)
            {
                const std::size_t iDof = element.nodes[iNode] - 1;  // 自由度索引（0-based）
                // 组装全局刚度矩阵三元组
                for (std::size_t jNode = 0; jNode < nodeCount; ++jNode)
                {
                    const std::size_t jDof = element.nodes[jNode] - 1;
                    kTriplet.emplace_back(iDof, jDof, eleKe(iNode, jNode));
                }
                // 累加全局荷载向量（Pw/PVy/PVz）
                Pw(iDof) += ele_FSet[0](iNode, 0);
                PVy(iDof) += ele_FSet[1](iNode, 0);
                PVz(iDof) += ele_FSet[2](iNode, 0);
            }         
        }
    }
    K.setFromTriplets (kTriplet.begin (), kTriplet.end ());

    //***********************************************添加约束，求解节点翘曲值、应力函数值***********************************************************
    // 添加约束(这里给第一个自由度添加固定约束0，采用乘大数法)
    const std::size_t constrainedDof = 0;                    // 第一个自由度添加固定约束0
    const Real        largeValue     = 1e18;                 // 大数
    Real              temp           = Pw (constrainedDof);  // 临时储存，便于还原Pe, 后面用于计算扭转常数J
    K.coeffRef (constrainedDof, constrainedDof) *= largeValue;
    Pw (constrainedDof) = 0.0;

    // 求解翘曲函数和剪切应力函数值
    Eigen::SimplicialLLT<EigenSparseR> solver;
    solver.compute (K);
    if (Eigen::Success != solver.info ())
    {
        std::cerr << "Decomposition failed !!!" << "fail type " << solver.info () << std::endl;
        assert (0);
    }
    const EigenVectorXR omega = solver.solve (Pw);
    const EigenVectorXR psi   = solver.solve (PVy);  // Vy应力函数
    const EigenVectorXR phi   = solver.solve (PVz);  // Vz应力函数

    //***********************************************计算扭转常数、剪心坐标、剪切常数***********************************************************
    // 求解截面扭转常数
    Pw (constrainedDof) = temp;
    sectionProperty.m_J = Iy + Iz - Pw.transpose () * omega;

    // 下面求剪心坐标（采用Trefftz's definition, 即与截面属性无关的剪心坐标）
    // 剪心坐标采用形心轴坐标系而非网格划分坐标系
    // 求扇形积
    Matrixd Iyzw (1, 2, 0.0);          // 剪心用：扇形积
    Matrixd integralTerm (3, 1, 0.0);  // 剪切系数用积分项
    for (const auto& type : existingTypes)
    {
        auto elementCalc = BeamSecElementFactory::Create (type);
        assert (elementCalc && "Failed to create element calculator");
        const auto& eleSet = meshes.GetElement (type);
        if (eleSet.empty()) continue;
        const std::size_t nodeCount = eleSet[0].nodes.size();    // 单元节点个数 
        Matrixd nodeCoords(nodeCount, 2);                        // 单元坐标
        Matrixd eleNodeOmega(nodeCount, 1);                      // 单元节点扇形积
        Matrixd eleNodePsi(nodeCount, 1);                        // 单元节点应力函数值
        Matrixd eleNodePhi(nodeCount, 1);

        for (const auto& element : eleSet)
        {
            nodeCoords = meshes.ExtractElementNodeCoords (element.nodes);
            
            // 计算扇形积
            for (std::size_t iNode = 0; iNode < nodeCount; ++iNode)
            {
                const std::size_t nodeIdx = element.nodes[iNode] - 1; // 1→0索引转换
                eleNodeOmega(iNode, 0) = omega(nodeIdx);
                eleNodePsi(iNode, 0) = psi(nodeIdx);   
                eleNodePhi(iNode, 0) = phi(nodeIdx);   
            }            
            Iyzw += elementCalc->ComputeSectorProduct(nodeCoords, eleNodeOmega);            

            // 计算剪切系数
            integralTerm += elementCalc->ComputeShearCoeff(nodeCoords, {Iy, Iz, Iyz}, mu, eleNodePsi, eleNodePhi);
        }
    }
    // 计算剪心坐标（Trefftz定义）
    temp = Iy * Iz - Iyz * Iyz;
    if (std::fabs (temp) < 1e-6)
    {
        SGFatal << "Denominator near zero for shear center!" << std::endl;
        return;
    }
    sectionProperty.m_yS = (Iyz * Iyzw (0, 0) - Iz * Iyzw (0, 1)) / temp;
    sectionProperty.m_zS = (Iy * Iyzw (0, 0) - Iyz * Iyzw (0, 1)) / temp;

    // 计算剪切系数
    const Real coef            = sectionProperty.m_area / (2.0 * (1.0 + mu) * (Iy * Iz - Iyz * Iyz) * 2.0 * (1.0 + mu) * temp);
    sectionProperty.m_alpha_y  = integralTerm (0, 0) * coef;
    sectionProperty.m_alpha_z  = integralTerm (1, 0) * coef;
    sectionProperty.m_alpha_yz = integralTerm (2, 0) * coef;
}